#include <pmm.h>
#include <list.h>
#include <string.h>
#include <slub_pmm.h>
#include<stdio.h>


free_area_t free_area;

#define free_list (free_area.free_list)
#define nr_free (free_area.nr_free)
#define buddy (free_area.buddy)
#define begin_page (free_area.begin_page)
struct kmem_cache kmem_cache_list[9]; // 9个 kmem_cache，管理从 8 字节到 2048 字节的对象缓存。

/* 初始化 free_list 和 nr_free */
static void default_init(void) {
    list_init(&free_list); // 初始化双向链表，存储空闲块
    nr_free = 0;           // 初始化空闲页数量为0
}

/* 打印 buddy 系统的节点容量 */
static void show_capacity(unsigned total_mem) {
    for (int i = 0; i < total_mem * 2 - 1; i++) {
        cprintf("%u ", buddy.node_capacity[i]); // 打印每个节点的容量
    }
    cprintf("\n\n");
}

/* 将大小转换为大于等于size的最接近2的幂 */
unsigned convert_2_pow_of_2(unsigned size) {
    unsigned temp = 0;
    unsigned back_up_size = size;

    while (back_up_size) {
        temp++;
        back_up_size /= 2; // 计算最接近的幂次
    }

    unsigned convert_2_pow = 1;
    for (int i = 0; i < temp - 1; i++)
        convert_2_pow *= 2;

    if (convert_2_pow != size)
        convert_2_pow *= 2; // 如果不相等，返回更大一档的幂

    return convert_2_pow; // 返回大于等于输入值的最接近的2的幂
}

/* 获取二叉树中节点的左孩子索引 */
unsigned left_child(unsigned index) {
    return index * 2 + 1; // 左孩子的公式：index * 2 + 1
}

/* 获取二叉树中节点的右孩子索引 */
unsigned right_child(unsigned index) {
    return index * 2 + 2; // 右孩子的公式：index * 2 + 2
}

/* 获取二叉树中节点的父节点索引 */
unsigned parent(unsigned index) {
    return (index - 1) / 2; // 父节点的公式： (index - 1) / 2
}

/* 初始化 buddy 系统 */
static void buddy_init(void) {
    nr_free = 0; // 初始化空闲页数量
}

/* 获取节点所在的层 */
static unsigned layer(unsigned index) {
    unsigned temp = 1;
    unsigned layer = 0;
    while (temp - 1 <= index) {
        layer++;
        temp *= 2;
    }
    return layer - 1;
}

/* 分配 buddy 系统中的页 */
static unsigned buddy_alloc_pages(size_t n) {
    assert(n > 0);

    if (n > buddy.node_capacity[0]) {
        return 0; // 如果请求超过 buddy 系统最大容量，返回 0
    }

    unsigned index = 0;
    unsigned offset = 0;
    unsigned pow_2_n = convert_2_pow_of_2(n); // 转换为最接近的 2 的幂

    int lor = 0; // 标识选择了左还是右孩子
    unsigned min = 0;

    // 在二叉树中寻找第一个能容纳请求块的节点
    for (int i = buddy.total_size; i != pow_2_n; i /= 2) {
        if (buddy.node_capacity[left_child(index)] <= buddy.node_capacity[right_child(index)]) {
            lor = 0;
            min = buddy.node_capacity[left_child(index)];
        } else {
            lor = 1;
            min = buddy.node_capacity[right_child(index)];
        }

        if (min >= pow_2_n) {
            index = lor ? right_child(index) : left_child(index);
        } else {
            index = !lor ? right_child(index) : left_child(index);
        }
    }

    buddy.node_capacity[index] = 0; // 找到后，占用整个块
    nr_free -= pow_2_n; // 更新空闲页数量

    /* 计算 offset */
    unsigned n_layer = layer(index);
    unsigned n_layer_pow = 1;
    for (int i = 0; i < n_layer; i++)
        n_layer_pow *= 2;
    offset = buddy.total_size / n_layer_pow * (index - n_layer_pow + 1);

    // 回溯更新节点容量，取左右孩子最大值
    while (index) {
        index = parent(index);
        buddy.node_capacity[index] = buddy.node_capacity[left_child(index)] > buddy.node_capacity[right_child(index)] 
                                    ? buddy.node_capacity[left_child(index)] 
                                    : buddy.node_capacity[right_child(index)];
    }

    return offset; // 返回分配的偏移
}

/* 初始化 buddy 系统的内存映射 */
void buddy_init_memmap(struct Page* base, unsigned n) {
    unsigned pow_n = convert_2_pow_of_2(n); // 确认请求大小是否为2的幂

    cprintf("start initializing\n");
    if (n < 1 || pow_n != n)
        return; // 请求不合法则返回

    nr_free = n; // 更新空闲页数量
    buddy.total_size = n;

    unsigned node_size = 2 * n; // 初始化节点的容量
    for (int i = 1; i < 2 * n; i++) {
        if (convert_2_pow_of_2(i) == i)
            node_size /= 2;

        buddy.node_capacity[i - 1] = node_size; // 初始化每个节点的容量
    }

    struct Page* p = base;
    for (; p != base + n; p++) {
        assert(PageReserved(p));
        p->flags = p->property = 0;
        set_page_ref(p, 0);
    }
    begin_page = base;

    // 链接 free_list
    free_list = base->page_link;
    struct Page* temp = le2page(&free_list, page_link);
    for (int i = 1; i < n; i++) {
        struct Page* p = (struct Page*)((char*)base + i * PGSIZE);
        temp->page_link.next = &p->page_link;
        p->page_link.prev = &temp->page_link;
        temp = p;
    }
    buddy_alloc_pages(1); // 保留页
    cprintf("already initialized\n");

    base->property = n;
    SetPageProperty(base);
    nr_free += n; // 更新空闲页
}

/* 释放 buddy 系统中的页 */
static void buddy_free_pages(size_t offset) {
    assert(offset >= 0); // 确保 offset 合法

    unsigned pow_2_n = 1;
    unsigned index = offset + buddy.total_size - 1;

    unsigned left, right;
    while (buddy.node_capacity[index]) {
        index = parent(index); // 查找父节点
        pow_2_n *= 2;          // 更新块大小
        if (!index) break;
    }

    buddy.node_capacity[index] = pow_2_n; // 恢复节点容量
    nr_free += pow_2_n; // 更新空闲页

    // 回溯更新父节点的容量
    while (index) {
        index = parent(index);
        pow_2_n *= 2;

        left = left_child(index);
        right = right_child(index);

        if (buddy.node_capacity[left] + buddy.node_capacity[right] == pow_2_n)
            buddy.node_capacity[index] = pow_2_n;
        else
            buddy.node_capacity[index] = buddy.node_capacity[left_child(index)] > buddy.node_capacity[right_child(index)]
                                        ? buddy.node_capacity[left_child(index)]
                                        : buddy.node_capacity[right_child(index)];
    }
}

/* 返回当前空闲的页数 */
static size_t buddy_nr_free_pages(void) {
    return nr_free;
}

/* 获取指定偏移的 Page */
struct Page* get_page(unsigned offset) {
    list_entry_t* le = &free_list;
    for (int i = 0; i < offset; i++) {
        le = list_next(le);
    }
    struct Page* p = le2page(le, page_link);
    return p;
}



/* 初始化 block_head，设置前后指针为空 */
void initialize_block_head(struct block_head* block) {
    block->next = NULL;
    block->prev = NULL;
}

/* 在链表尾部插入新的块 */
void insert(struct block_head* list_head, struct block_head* inserted_node) {
    struct block_head* p = list_head;
    while (p->next) {
        p = p->next; // 遍历到链表末尾
    }
    p->next = inserted_node;  // 插入新的节点
    inserted_node->prev = p;  // 更新前后指针
}

/* 从链表中删除指定节点 */
void del(struct block_head* list_head, struct block_head* node_2_del) {

    struct block_head* p = list_head;
    while (p->next) {
        if (p == node_2_del) {
            struct block_head* pre = node_2_del->prev;
            struct block_head* next = node_2_del->next;
            if (pre) {
                pre->next = next; // 更新前节点的 next 指针
            }
            if (next) {
                next->prev = pre; // 更新后节点的 prev 指针
            }
            return;
        }
        p = p->next; // 遍历链表
    }
}


/* 初始化 SLUB 分配器的内存映射 */
void slub_init_memmap(struct Page* base, unsigned n) {
    // 首先初始化 buddy 系统
    buddy_init_memmap(base, n);

    // 初始化 kmem_cache_list 数组，每个缓存管理不同大小的对象
    unsigned size = 8;
    for (int i = 0; i < 9; i++) {
        kmem_cache_list[i].object_size = size; // 设置缓存的对象大小
        initialize_block_head(&kmem_cache_list[i].object_seller.object_list_head); // 初始化每个缓存的 object_seller 链表
        kmem_cache_list[i].node_storage.partial_heads = NULL; // 初始化部分填充 slab 页的链表
        kmem_cache_list[i].node_storage.full_heads = NULL;    // 初始化已满 slab 页的链表
        size *= 2; // 每个缓存大小是前一个的两倍
    }
}

/* 查找 object_seller 链表中的空闲节点 */
int free_node_index(struct block_head* head) {
    if (head->next) {
        return 1; // 如果链表中有节点返回1
    }
    return -1; // 否则返回-1
}

/* 打印链表中的节点地址 */
void show_list(struct block_head* head) {
    struct block_head* p = head;
    while (p) {
        cprintf("%x\n", p); // 打印每个节点的地址
        p = p->next;
    }
}

/* 获取 kmem_cache_node 中部分填充的 slab 页 */
struct Page* get_partial_storage(struct kmem_cache_node* storage) {
    struct Page* p = storage->partial_heads;
    if (p != NULL && p->storage_head != NULL) {
        return p; // 返回部分填充的 slab 页
    }
    return NULL; // 没有可用的 slab 页时返回 NULL
}

/* 初始化 SLUB 分配器 */
static void slub_init(void) {
    nr_free = 0; // 初始化空闲页数量
}

/* 根据对象大小计算对应的 kmem_cache_list 索引 */
unsigned kmalloc_caches_index(size_t n) {
    unsigned temp = 0;
    while (n) {
        n /= 2; // 将对象大小除以2，找到合适的缓存
        temp++;
    }
    return temp - 4; // 返回对象大小对应的缓存索引
}

unsigned test_used_recorder = 0;


/* 分配 SLUB 缓存中的页面 */
static struct block_head* slub_alloc_pages(size_t n) {
    unsigned pow_n = convert_2_pow_of_2(n); // 将 n 转换为 2 的幂
    unsigned object_nums = 1;
    
    if (pow_n > 2048) { // 如果请求大小超过 2048 字节，调整对象数量
        object_nums = pow_n / 2048 * 2048 == pow_n ? pow_n / 2048 : pow_n / 2048 + 1;
        pow_n = 2048;
    }
    unsigned index = kmalloc_caches_index(pow_n); // 获取 kmem_cache_list 的索引

    // 查找 object_seller 链表中的空闲块
    int free_node_pos = free_node_index(&kmem_cache_list[index].object_seller.object_list_head);
    struct Page* partial_storage = get_partial_storage(&kmem_cache_list[index].node_storage);

    if (free_node_pos != -1) {
        struct block_head* p = kmem_cache_list[index].object_seller.object_list_head.next;
        kmem_cache_list[index].object_seller.object_list_head.next = p->next;

        if (p->next) {
            p->next->prev = &kmem_cache_list[index].object_seller.object_list_head;
        } else {
            // 如果 object_seller 中没有更多空闲块，更新缓存
            if (kmem_cache_list[index].node_storage.full_heads) {
                list_entry_t* le = &kmem_cache_list[index].node_storage.full_heads->page_link;
                for (int i = 1; i < kmem_cache_list[index].node_storage.full_num; i++) {
                    le = list_next(le);
                }
                struct Page* temp = le2page(le, page_link);
                temp->page_link.next = &kmem_cache_list[index].base->page_link;
                kmem_cache_list[index].base->page_link.prev = temp->page_link.next;
            } else {
                kmem_cache_list[index].node_storage.full_heads = kmem_cache_list[index].base;
            }
            kmem_cache_list[index].base = NULL;
            kmem_cache_list[index].node_storage.full_num++; // 增加已满的仓库数量
        }
        return p;
    } else if (partial_storage != NULL) {
        // 如果有部分填充的 slab，分配对象并更新链表
        struct block_head* object = partial_storage->storage_head;
        if (object->next) {
            object->next->prev = NULL;
            partial_storage->storage_head = object->next;
            return object;
        } else {
            // 如果 slab 已经满了，转移到 full_heads 链表
            partial_storage->storage_head = NULL;
            kmem_cache_list[index].node_storage.full_num++;
            kmem_cache_list[index].node_storage.partial_num--;

            list_del(&partial_storage->page_link);
            list_add(&kmem_cache_list[index].node_storage.full_heads->page_link, &partial_storage->page_link);
            return object;
        }
    } else {
        // 没有空闲块和部分 slab 时，从 buddy 系统分配新页
        test_used_recorder++;
        unsigned page_num = n / PGSIZE * PGSIZE == n ? n / PGSIZE : n / PGSIZE + 1;
        unsigned offset = buddy_alloc_pages(page_num);
        struct Page* base_page = get_page(offset); // 获取分配的新页

        kmem_cache_list[index].object_seller.object_list_head.next = NULL;
        kmem_cache_list[index].object_seller.object_list_head.prev = NULL;

        kmem_cache_list[index].base = base_page;
        kmem_cache_list[index].object_seller.object_list_head = *((struct block_head*)((char*)base_page + pow_n));

        struct block_head* p = &kmem_cache_list[index].object_seller.object_list_head;
        for (int i = 1; i < PGSIZE / pow_n - 1; i++) {
            struct block_head* temp = (struct block_head*)((char*)base_page + pow_n * (i + 2)); // 划分对象
            initialize_block_head(temp);
            temp->size = index;
            p->next = temp;
            temp->prev = p;
            p = temp;
        }
        
        struct block_head* obj = kmem_cache_list[index].object_seller.object_list_head.next;
        kmem_cache_list[index].object_seller.object_list_head.next = obj->next;
        if (obj->next) {
            obj->next->prev = &kmem_cache_list[index].object_seller.object_list_head;
        }
                // 更新 object_seller 链表中的头节点
      cprintf("\n满页啦!!\n");
       cprintf("基地址:%x   ADDRESS:%x\n", base_page, obj); // 打印分配的基地址和分配的对象地址
        return obj; // 返回分配的对象地址
    }
}

/* 获取存储满块的 Page */
struct Page* get_storage_full(struct block_head* object) {
    unsigned index = object->size; // 获取对象的大小
    struct Page* p;

    list_entry_t* le = &kmem_cache_list[index].node_storage.full_heads->page_link;
    for (int i = 0; i < kmem_cache_list[index].node_storage.full_num; i++) {
        p = le2page(le, page_link);
        if ((char*)object - (char*)p < PGSIZE && (char*)object - (char*)p > 0) {
            return p; // 返回对象所在的 slab 页
        }
        le = list_next(le);
    }
    return NULL; // 没有找到对应的 slab 页，返回 NULL
}

/* 获取部分填充块的 Page */
struct Page* get_storage_partial(struct block_head* object) {
    unsigned index = object->size; // 获取对象的大小
    struct Page* p;

    list_entry_t* le = &kmem_cache_list[index].node_storage.partial_heads->page_link;
    for (int i = 0; i < kmem_cache_list[index].node_storage.partial_num; i++) {
        p = le2page(le, page_link);
        if ((char*)object - (char*)p < PGSIZE && (char*)object - (char*)p > 0) {
            return p; // 返回对象所在的部分填充 slab 页
        }
        le = list_next(le);
    }
    return NULL; // 没有找到对应的部分 slab 页，返回 NULL
}

/* 释放 SLUB 分配的页面 */
static void slub_free_pages(struct block_head* object) {
    unsigned index = object->size; // 获取对象的大小
    struct Page* storage_full = get_storage_full(object); // 获取满的 slab 页
    struct Page* storage_partial = get_storage_partial(object); // 获取部分填充的 slab 页

   /* if (storage_full)
        cprintf("满页:%x\n", storage_full); // 如果有满的 slab 页，打印其地址
    if (storage_partial)
        cprintf("部分slab页:%x\n", storage_partial); // 如果有部分 slab 页，打印其地址
*/
    // 检查对象是否属于当前的 base_page
    if ((char*)object - (char*)kmem_cache_list[index].base < PGSIZE && (char*)object - (char*)kmem_cache_list[index].base >= 0) {
        // 当前释放的对象属于当前的缓存，重新链接到 object_seller 链表
        struct block_head* p = &kmem_cache_list[index].object_seller.object_list_head;
        while (p < object) {
            if (p->next == NULL) { // 如果 object 是最后一个，直接插入
                p->next = object;
                object->prev = p;
                return;
            }
            p = p->next;
        }
        // 插入 object
        p->prev->next = object;
        object->prev = p->prev;
        object->next = p;
        p->prev = object;
        return;
    } 
    // 如果对象属于满的 slab 页
    else if (storage_full != NULL) {
        if (kmem_cache_list[index].node_storage.full_heads == storage_full) {
            // 如果是满的第一个节点，更新 full_heads
            struct Page* full_test_head = le2page(list_next(&kmem_cache_list[index].node_storage.full_heads->page_link), page_link);
            full_test_head->page_link.prev = NULL;
            kmem_cache_list[index].node_storage.full_heads = full_test_head;
            kmem_cache_list[index].node_storage.full_num--; // 减少满的 slab 页数量
        } else {
            // 否则从 full_heads 中删除
            list_entry_t* le = &kmem_cache_list[index].node_storage.full_heads->page_link;
            while (le2page(le, page_link) != storage_full) {
                le = list_next(le);
            }
            list_del(le);
            kmem_cache_list[index].node_storage.full_num--; // 减少满的 slab 页数量
        }

        // 将 slab 页转移到 partial_heads
        if (kmem_cache_list[index].node_storage.partial_heads == NULL) {
            kmem_cache_list[index].node_storage.partial_heads = storage_full;
        } else {
            list_add(&kmem_cache_list[index].node_storage.partial_heads->page_link, &storage_full->page_link);
        }
        kmem_cache_list[index].node_storage.partial_num++; // 增加部分 slab 页数量

        storage_full->storage_head = object; // 重新链接对象到 partial_heads
        storage_full->storage_head->next = NULL;
        storage_full->storage_head->prev = NULL;
    }
    // 如果对象属于部分 slab 页
    else if (storage_partial != NULL) {
        // 插入到部分 slab 页的链表
        if (object < storage_partial->storage_head) {
            storage_partial->storage_head->prev = object;
            object->next = storage_partial->storage_head;
            storage_partial->storage_head = object;
            object->prev = NULL;
            return;
        } else {
            struct block_head* p = storage_partial->storage_head;
            while (p < object) {
                if (p->next == NULL) {
                    p->next = object;
                    object->prev = p;
                    return;
                }
                p = p->next;
            }
            p->prev->next = object;
            object->prev = p->prev;
            object->next = p;
            p->prev = object;
            return;
        }
    }
}

/* 返回当前 SLUB 空闲页的数量 */
static size_t slub_nr_free_pages(void) {
    return nr_free;
}

/* 打印 slab 链表中的页面和对象 */
void print_pagelist(struct Page* head, unsigned num) {
    struct Page* p = head;

    for (int i = 0; i < num; i++) {
        cprintf("该slab 页中对象的链表 %x slab\n", p);
        show_list(p->storage_head); // 打印每个 slab 页中对象的链表
        cprintf("--------------------------------------------------\n");
        p = le2page(list_next(&p->page_link), page_link);
    }
}

/* 测试 SLUB 分配器的基本功能 */
static void slub_basic_check(void) 
{
	cprintf("--------------------------------------------------\n\n");
	cprintf("slub检查开始!!\n");
    struct block_head *all_512s[20];
	cprintf("分配 20 个 512 字节的对象:\n");
    // 分配 10 个 512 字节的对象
    for (int i = 0; i < 20; i++) {
        all_512s[i] = slub_alloc_pages(512);
        cprintf("ADDRESS: %x\n", all_512s[i]);
    }

    // 打印已满的 slab 页
    struct Page* full_test_head;
    list_entry_t* le = &kmem_cache_list[kmalloc_caches_index(512)].node_storage.full_heads->page_link;
    for (int i = 0; i < kmem_cache_list[kmalloc_caches_index(512)].node_storage.full_num; i++) {
        full_test_head = le2page(le, page_link);
        cprintf("满页: %x\n", full_test_head);
        le = list_next(le);
    }
    
    cprintf("--------------------------------------------------\n");
	cprintf("释放部分分配的对象:\n");
    // 释放部分分配的对象
    slub_free_pages(all_512s[0]);
    slub_free_pages(all_512s[4]);

cprintf("打印部分填充的 slab 页:\n");
    // 打印部分填充的 slab 页
    print_pagelist(kmem_cache_list[kmalloc_caches_index(512)].node_storage.partial_heads, 
                   kmem_cache_list[kmalloc_caches_index(512)].node_storage.partial_num);


cprintf("继续分配 2 个 512 字节的对象:\n");
    // 继续分配 2 个 512 字节的对象
    for (int i = 8; i < 10; i++) {
        all_512s[i] = slub_alloc_pages(512);
        cprintf("ADDRESS: %x\n", all_512s[i]);
    }

    // 测试完成标志
    cprintf("--------------------------------------------------\n");
    cprintf("Slub检查完毕!!\n");
}

/* 定义 SLUB 分配器的 pmm_manager */
const struct pmm_manager slub_pmm_manager = {
    .name = "slub_pmm_manager",         // 分配器的名字
    .init = slub_init,                  // 分配器的初始化函数
    .init_memmap = slub_init_memmap,    // 初始化内存映射
    .alloc_pages = slub_alloc_pages,    // 分配页面
    .free_pages = slub_free_pages,      // 释放页面
    .nr_free_pages = slub_nr_free_pages, // 返回空闲页面数量
    .check = slub_basic_check,          // 基本功能测试
};
